// Helper for Tcp communication.

#include "TcpHelper.h"
#include "BitHelper.h"
#include <ws2tcpip.h>

// Construct TcpHelper with given address and port.
TcpHelper::TcpHelper(const string& address, const string& port)
{
	m_connected = false;
	m_address = address;
	m_port = port;
	m_socket = INVALID_SOCKET;
}

// Finalize TcpHelper.
TcpHelper::~TcpHelper()
{
	if (m_connected) TryDisconnect();
}

// Try to connect to the server.
bool TcpHelper::TryConnect()
{
	if (m_connected) return LogManager::DisplayWarning(L"You are already connected to the server !\n");

	PADDRINFOW address = NULL;

	// initialize Winsock, resolve the server address and port, then attempt to connect to an address list until one succeeds
	if (TryInitializeWinsock() && TryResolveAddressInfo(address) && TryConnectAddress(address))
	{
		m_connected = true;
		return LogManager::DisplayInformation(L"You are now connected to the server.\n");
	}

	return false;
}

// Try to disconnect from the server.
bool TcpHelper::TryDisconnect()
{
	if (!m_connected) return LogManager::DisplayWarning(L"You are already disconnected from the server !\n");

	int iResult = shutdown(m_socket, SD_SEND);
	m_connected = false;
	if (iResult == SOCKET_ERROR) return HandleErrorAndClose("The socket could not be shut down");

	closesocket(m_socket);
	WSACleanup();
	return LogManager::DisplayInformation(L"You are now disconnected from the server.\n");
}

// Try to send a buffer to the server.
bool TcpHelper::TrySend(vector<char>& buffer)
{
	int iResult = send(m_socket, &buffer[0], static_cast<int>(buffer.size()), 0);
	if (iResult == SOCKET_ERROR) return HandleErrorAndDisconnect("Data could not be sent to the server");
	return true;
}

// Read a value of a generic type from the server.
template<typename T>
bool TcpHelper::TryRead(T& result)
{
	return TryRead(result, sizeof(T));
}

// Try to read a string from the server.
template<>
bool TcpHelper::TryRead(string& result)
{
	int length;
	return TryRead7BitEncodedInt(length) && TryRead(result, length);
}

// Try to initialize Winsock.
bool TcpHelper::TryInitializeWinsock()
{
	WSADATA wsaData;
	int iResult = WSAStartup(MAKEWORD(2,2), &wsaData);
	if (iResult == 0) return true;
	return LogManager::DisplayWarning(L"Winsock could not be initialized (error: %d).\n", iResult);
}

// Try to resolve the server address and port.
bool TcpHelper::TryResolveAddressInfo(PADDRINFOW& result)
{
	ADDRINFOW hints;
	ZeroMemory(&hints, sizeof(hints));
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_protocol = IPPROTO_TCP;

	int iResult = GetAddrInfoW(_W(m_address).c_str(), _W(m_port).c_str(), &hints, &result);
	if (iResult != 0) return HandleError("Server address and port could not be resolved", iResult);
	return true;
}

// Attempt to connect to an address list until one succeeds.
bool TcpHelper::TryConnectAddress(const PADDRINFOW& address)
{
	PADDRINFOW ptr = NULL;
	for(ptr = address; ptr != NULL; ptr = ptr->ai_next)
	{
		// try to create a socket then to use it to connect to server
		if (TryCreateSocket(ptr) && TryConnectServer(ptr)) break;
	}

	// frees address information dynamically allocated
	FreeAddrInfoW(address);

	// was unable to connect to any server
	if (m_socket == INVALID_SOCKET) return HandleError("Unable to connect to server");
	return true;
}

// Try to create a socket for connecting to server.
bool TcpHelper::TryCreateSocket(const PADDRINFOW& address)
{
	LogManager::DisplayInformation(L"TryCreateSocket... ai_family : %d /z ai_socktype : %d ; ai_protocol : %d.\n", address->ai_family, address->ai_socktype, address->ai_protocol);
	m_socket = socket(address->ai_family, address->ai_socktype, address->ai_protocol);
	int i = 1;
	setsockopt(m_socket, IPPROTO_TCP, TCP_NODELAY, (const char*)&i, INT_SIZE);
	if (m_socket == INVALID_SOCKET) return HandleError("The socket creation failed");
	return true;
}

// Try to connect to server.
bool TcpHelper::TryConnectServer(const PADDRINFOW& address)
{
	LogManager::DisplayInformation(L"TryConnectServer... ai_family : %d ; ai_socktype : %d ; ai_protocol : %d.\n", address->ai_family, address->ai_socktype, address->ai_protocol);
	int iResult = connect(m_socket, address->ai_addr, (int)address->ai_addrlen);
	if (iResult == SOCKET_ERROR)
	{
		HandleErrorAndClose("Connection to the server failed");
		m_socket = INVALID_SOCKET;
		return false;
	}
	return true;
}

// Read a value of a generic type from the server.
template<typename T>
bool TcpHelper::TryRead(T& result, int length)
{
	vector<char> buffer(length);
	int iResult = recv(m_socket, &buffer[0], length, 0);
	if (iResult > 0)
	{
		result = BitHelper::Convert<T>(buffer);
		return true;
	}
	if (iResult < 0) return HandleErrorAndDisconnect("A value could not be read from the server");
	return HandleErrorAndDisconnect("The server did not respond in time");
}

// Dummy helper to prevent LNK errors
void TcpHelperLnkHelper()
{
	TcpHelper tcpHelper("", "");
	byte dummyByte;
	bool dummyBool;
	int dummyInt;
	unsigned int dummyUnsignedInt;
	double dummyDouble;
	string dummyString;
	tcpHelper.TryRead(dummyByte);
	tcpHelper.TryRead(dummyBool);
	tcpHelper.TryRead(dummyInt);
	tcpHelper.TryRead(dummyUnsignedInt);
	tcpHelper.TryRead(dummyDouble);
	tcpHelper.TryRead(dummyString);
}

// Try to reads in a 32-bit integer in compressed format.
bool TcpHelper::TryRead7BitEncodedInt(int& result)
{
	result = 0;
	int bitIndex = 0;
	byte currentByte;

	while (bitIndex != 35)
	{
		if (!TryRead(currentByte)) return false;
		result |= ((int)currentByte & (int)SCHAR_MAX) << bitIndex;
		bitIndex += 7;

		if (((int)currentByte & 128) == 0)
			return true;
	}
	return HandleErrorAndDisconnect("A 7 bit encoded integer could not be read from the server");
}

// Print the error message and clean the managed resources.
bool TcpHelper::HandleError(const string& message)
{
	WSACleanup();
	return LogManager::DisplayWarning(L"%s (address : %s ; port : %s).\n", _W(message).c_str(), _W(m_address).c_str(), _W(m_port).c_str());
}

// Print the error message, close the socket and clean the managed resources.
bool TcpHelper::HandleError(const string& message, const int& error)
{
	WSACleanup();
	return LogManager::DisplayWarning(L"%s (address : %s ; port : %s ; error: %d).\n", _W(message).c_str(), _W(m_address).c_str(), _W(m_port).c_str(), error);
}

// Close the socket, then print the error message and clean the managed resources.
bool TcpHelper::HandleErrorAndClose(const string& message)
{
	int error = WSAGetLastError();
	closesocket(m_socket);
	return HandleError(message, error);
}

// Print the error message then try to disconnect from the server.
bool TcpHelper::HandleErrorAndDisconnect(const string& message)
{
	int error = WSAGetLastError();
	LogManager::DisplayWarning(L"%s (address : %s ; port : %s ; error: %d).\n", _W(message).c_str(), _W(m_address).c_str(), _W(m_port).c_str(), error);
	return TryDisconnect() && false;
}